#ifndef MYAMG_MATRIXFEATURES_CUH
#define MYAMG_MATRIXFEATURES_CUH
#include <limits>
#include <Data/MatrixFeaturesKernel.cuh>
#include <cassert>
#include <unordered_map>

namespace myamg {
    class MatrixFeatures {
    public:
        /*
         * Basic Features
        */
        // The shape of the matrix
        size_t row;
        size_t col;
        // NNZ count of the matrix
        size_t nnz;
        // The number of positive diagonal elements
        size_t positive_diagonal_count = 0;
        // The number of negative diagonal elements
        size_t negative_diagonal_count = 0;
        // The number of positive non - diagonal elements
        size_t positive_nondiagonal_count = 0;
        // The number of negative non - diagonal elements
        size_t negative_nondiagonal_count = 0;
        // The number of diangonally dominant rows
        // Definition of diagonally dominant : | diagonal element| >= sum(| non - diagonal elements | )
        size_t diagonally_dominant_row_count = 0;
        // The number of non - diagonally dominant rows
        // Definition of diagonally dominant : | diagonal element| >= sum(| non - diagonal elements | )
        size_t non_diagonally_dominant_row_count = 0;
        // (diagonally dominant row count)/(total row count)
        double diagonally_dominant_row_percentage = 0.0;
        // (non-diagonally dominant row count)/(total row count)
        double non_diagonally_dominant_row_percentage = 0.0;

        /*
         * Multiscale Features
         */

        // minimum row nnz count of the matrix
        size_t min_row_nnz_count = std::numeric_limits<size_t>::max();
        // maximum row nnz count of the matrix
        size_t max_row_nnz_count = 0;
        // minimum | diagonal element |
        double min_abs_diagonal = std::numeric_limits<double>::max();
        // maximum | diagonal element |
        double max_abs_diagonal = 0.0;
        // minimum | non - diagonal element |
        double min_abs_non_diagonal = std::numeric_limits<double>::max();
        // maximum | non - diagonal element |
        double max_abs_non_diagonal = 0.0;
        // min(max(| non - diagonal element | ) / min(| non - diagonal element | ))
        double min_frac_nondiag_max_vs_min = std::numeric_limits<double>::max();
        // max(max(| non - diagonal element | ) / min(| non - diagonal element | ))
        double max_frac_nondiag_max_vs_min = 0.0;

        void clear() {
            row = 0;
            col = 0;
            nnz = 0;
            positive_diagonal_count = 0;
            negative_diagonal_count = 0;
            positive_nondiagonal_count = 0;
            negative_nondiagonal_count = 0;
            diagonally_dominant_row_count = 0;
            non_diagonally_dominant_row_count = 0;
            diagonally_dominant_row_percentage = 0.0;
            non_diagonally_dominant_row_percentage = 0.0;
            min_row_nnz_count = std::numeric_limits<size_t>::max();
            max_row_nnz_count = 0;
            min_abs_diagonal = std::numeric_limits<double>::max();
            max_abs_diagonal = 0.0;
            min_abs_non_diagonal = std::numeric_limits<double>::max();
            max_abs_non_diagonal = 0.0;
            min_frac_nondiag_max_vs_min = std::numeric_limits<double>::max();
            max_frac_nondiag_max_vs_min = 0.0;
        }

        template<typename Integer, typename Real, typename Allocator>
        void compute(const CSRMatrix<Integer, Real, Allocator> &matrix) {
            clear();
            row = matrix.numRows();
            col = matrix.numCols();
            nnz = matrix.numNonzeros();

            auto *rowptr = new Integer[row + 1];
            auto *colind = new Integer[nnz];
            auto *values = new Real[nnz];

            if constexpr (AllocatorType<Allocator>::value == MemoryType::Host) {
                std::copy(matrix.rowPtr(), matrix.rowPtr() + row + 1, rowptr);
                std::copy(matrix.colInd(), matrix.colInd() + nnz, colind);
                std::copy(matrix.values(), matrix.values() + nnz, values);
            } else if constexpr (AllocatorType<Allocator>::value == MemoryType::Cuda) {
                cudaMemcpy(rowptr, matrix.rowPtr(), (row + 1) * sizeof(Integer), cudaMemcpyDeviceToHost);
                cudaMemcpy(colind, matrix.colIndex(), nnz * sizeof(Integer), cudaMemcpyDeviceToHost);
                cudaMemcpy(values, matrix.values(), nnz * sizeof(Real), cudaMemcpyDeviceToHost);
            } else {
                std::cerr << "MatrixFeatures::compute(): Unknown allocator type." << std::endl;
                delete[] rowptr;
                delete[] colind;
                delete[] values;
                return;
            }

            for (auto i = 0; i < row; ++i) {
                auto row_nnz_count = rowptr[i + 1] - rowptr[i];
                if (row_nnz_count < min_row_nnz_count) {
                    min_row_nnz_count = row_nnz_count;
                }
                if (row_nnz_count > max_row_nnz_count) {
                    max_row_nnz_count = row_nnz_count;
                }
                Real diag = 0;
                Real sum_nondiag = 0;
                Real min_nondiag = std::numeric_limits<Real>::max();
                Real max_nondiag = 0;
                size_t n_nondiag = 0;
                for (auto j = rowptr[i]; j < rowptr[i + 1]; ++j) {
                    auto jcol = colind[j];
                    auto val = values[j];
                    auto absVal = std::abs(val);
                    if (i == jcol) {
                        diag = absVal;
                        if (absVal > max_abs_diagonal) {
                            max_abs_diagonal = absVal;
                        }
                        if (absVal < min_abs_diagonal) {
                            min_abs_diagonal = absVal;
                        }
                        if (val > 0) {
                            ++positive_diagonal_count;
                        } else if (val < 0) {
                            ++negative_diagonal_count;
                        }
                    } else {
                        ++n_nondiag;
                        sum_nondiag += absVal;
                        if (absVal > max_abs_non_diagonal) {
                            max_abs_non_diagonal = absVal;
                        }
                        if (absVal < min_abs_non_diagonal) {
                            min_abs_non_diagonal = absVal;
                        }
                        if (absVal > max_nondiag) {
                            max_nondiag = absVal;
                        }
                        if (absVal < min_nondiag) {
                            min_nondiag = absVal;
                        }
                        if (val > 0) {
                            ++positive_nondiagonal_count;
                        } else if (val < 0) {
                            ++negative_nondiagonal_count;
                        }
                    }
                }
                if (diag >= sum_nondiag) {
                    ++diagonally_dominant_row_count;
                } else {
                    ++non_diagonally_dominant_row_count;
                }
                if (n_nondiag > 0) {
                    auto frac = max_nondiag / min_nondiag;
                    if (frac < min_frac_nondiag_max_vs_min) {
                        min_frac_nondiag_max_vs_min = frac;
                    }
                    if (frac > max_frac_nondiag_max_vs_min) {
                        max_frac_nondiag_max_vs_min = frac;
                    }
                }
            }
            assert((diagonally_dominant_row_count + non_diagonally_dominant_row_count) == row);
            diagonally_dominant_row_percentage =
                    static_cast<double>(diagonally_dominant_row_count) / static_cast<double>(row);
            non_diagonally_dominant_row_percentage =
                    static_cast<double>(non_diagonally_dominant_row_count) / static_cast<double>(row);
            delete[] rowptr;
            delete[] colind;
            delete[] values;
        }

        std::vector<std::pair<std::string, std::string> > toMap(const std::string& name = "") const {
            std::vector<std::pair<std::string, std::string> > featureMap;

            // Always add name field first (empty string if not set)
            featureMap.emplace_back("name", name);

            // Basic features
            featureMap.emplace_back("row", std::to_string(row));
            featureMap.emplace_back("col", std::to_string(col));
            featureMap.emplace_back("nnz", std::to_string(nnz));
            featureMap.emplace_back("positive_diagonal_count", std::to_string(positive_diagonal_count));
            featureMap.emplace_back("negative_diagonal_count", std::to_string(negative_diagonal_count));
            featureMap.emplace_back("positive_nondiagonal_count", std::to_string(positive_nondiagonal_count));
            featureMap.emplace_back("negative_nondiagonal_count", std::to_string(negative_nondiagonal_count));
            featureMap.emplace_back("diagonally_dominant_row_count", std::to_string(diagonally_dominant_row_count));
            featureMap.emplace_back("non_diagonally_dominant_row_count",
                                    std::to_string(non_diagonally_dominant_row_count));
            featureMap.emplace_back("diagonally_dominant_row_percentage",
                                    std::to_string(diagonally_dominant_row_percentage));
            featureMap.emplace_back("non_diagonally_dominant_row_percentage",
                                    std::to_string(non_diagonally_dominant_row_percentage));

            // Multiscale features
            featureMap.emplace_back("min_row_nnz_count", std::to_string(min_row_nnz_count));
            featureMap.emplace_back("max_row_nnz_count", std::to_string(max_row_nnz_count));
            featureMap.emplace_back("min_abs_diagonal", std::to_string(min_abs_diagonal));
            featureMap.emplace_back("max_abs_diagonal", std::to_string(max_abs_diagonal));
            featureMap.emplace_back("min_abs_non_diagonal", std::to_string(min_abs_non_diagonal));
            featureMap.emplace_back("max_abs_non_diagonal", std::to_string(max_abs_non_diagonal));
            featureMap.emplace_back("min_frac_nondiag_max_vs_min", std::to_string(min_frac_nondiag_max_vs_min));
            featureMap.emplace_back("max_frac_nondiag_max_vs_min", std::to_string(max_frac_nondiag_max_vs_min));

            return featureMap;
        }
    };
}

#endif //MYAMG_MATRIXFEATURES_CUH
